// ROS官方相机标定工具
// 功能：使用ROS官方标定工具对松灵小车的color相机进行标定

#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/CameraInfo.h>
#include <cv_bridge/cv_bridge.h>
#include <opencv2/opencv.hpp>
#include <opencv2/calib3d.hpp>
#include <iostream>
#include <fstream>
#include <vector>

class ROSCameraCalibration {
private:
    ros::NodeHandle nh_;
    ros::Subscriber image_sub_;
    ros::Subscriber camera_info_sub_;
    
    // 标定参数
    std::vector<std::vector<cv::Point3f>> object_points_;
    std::vector<std::vector<cv::Point2f>> image_points_;
    cv::Mat camera_matrix_;
    cv::Mat distortion_coeffs_;
    std::vector<cv::Mat> rvecs_, tvecs_;
    
    // 标定板参数
    cv::Size board_size_;
    float square_size_;
    cv::Size image_size_;
    
    // 标定状态
    bool calibration_done_;
    int images_captured_;
    int required_images_;
    
    std::string output_file_;

public:
    ROSCameraCalibration() : 
        board_size_(9, 6),  // 9x6 内角点
        square_size_(0.025), // 25mm
        calibration_done_(false),
        images_captured_(0),
        required_images_(20),
        output_file_("ros_camera_calibration_result.yaml") {
        
        image_sub_ = nh_.subscribe("/camera/color/image_raw", 1, 
                                 &ROSCameraCalibration::imageCallback, this);
        camera_info_sub_ = nh_.subscribe("/camera/color/camera_info", 1,
                                       &ROSCameraCalibration::cameraInfoCallback, this);
        
        ROS_INFO("ROS Camera Calibration initialized");
        ROS_INFO("Board size: %dx%d, Square size: %.3f m", 
                board_size_.width, board_size_.height, square_size_);
        ROS_INFO("Required images: %d", required_images_);
        ROS_INFO("Press 'c' to capture image, 'q' to quit, 's' to start calibration");
    }

    void imageCallback(const sensor_msgs::Image::ConstPtr& msg) {
        if (calibration_done_) return;
        
        cv_bridge::CvImagePtr cv_ptr;
        try {
            cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
            return;
        }
        
        cv::Mat image = cv_ptr->image;
        image_size_ = image.size();
        
        // 查找棋盘格角点
        std::vector<cv::Point2f> corners;
        bool found = cv::findChessboardCorners(image, board_size_, corners,
                                             cv::CALIB_CB_ADAPTIVE_THRESH | 
                                             cv::CALIB_CB_NORMALIZE_IMAGE);
        
        if (found) {
            // 亚像素精度优化
            cv::cornerSubPix(image, corners, cv::Size(11, 11), cv::Size(-1, -1),
                           cv::TermCriteria(cv::TermCriteria::EPS + cv::TermCriteria::COUNT, 30, 0.1));
            
            // 绘制角点
            cv::drawChessboardCorners(image, board_size_, corners, found);
        }
        
        // 显示图像和状态
        std::string status = "Images captured: " + std::to_string(images_captured_) + 
                           "/" + std::to_string(required_images_);
        cv::putText(image, status, cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 
                   0.7, cv::Scalar(0, 255, 0), 2);
        
        cv::imshow("ROS Camera Calibration", image);
        char key = cv::waitKey(1);
        
        if (key == 'c' && found && images_captured_ < required_images_) {
            // 捕获标定图像
            captureCalibrationImage(corners);
        } else if (key == 's' && images_captured_ >= 10) {
            // 开始标定
            performCalibration();
        } else if (key == 'q') {
            ros::shutdown();
        }
    }
    
    void cameraInfoCallback(const sensor_msgs::CameraInfo::ConstPtr& msg) {
        // 获取初始相机参数（如果可用）
        if (!calibration_done_) {
            ROS_INFO("Received initial camera info");
        }
    }
    
    void captureCalibrationImage(const std::vector<cv::Point2f>& corners) {
        // 生成3D对象点
        std::vector<cv::Point3f> object_points;
        for (int i = 0; i < board_size_.height; i++) {
            for (int j = 0; j < board_size_.width; j++) {
                object_points.push_back(cv::Point3f(j * square_size_, i * square_size_, 0));
            }
        }
        
        object_points_.push_back(object_points);
        image_points_.push_back(corners);
        images_captured_++;
        
        ROS_INFO("Captured image %d/%d", images_captured_, required_images_);
    }
    
    void performCalibration() {
        if (images_captured_ < 10) {
            ROS_WARN("Not enough images for calibration. Need at least 10, have %d", images_captured_);
            return;
        }
        
        ROS_INFO("Starting camera calibration with %d images...", images_captured_);
        
        // 执行相机标定
        std::vector<cv::Mat> rvecs, tvecs;
        double rms = cv::calibrateCamera(object_points_, image_points_, image_size_,
                                       camera_matrix_, distortion_coeffs_, rvecs, tvecs);
        
        ROS_INFO("Calibration completed!");
        ROS_INFO("RMS reprojection error: %.3f pixels", rms);
        
        // 保存标定结果
        saveCalibrationResults(rms);
        
        calibration_done_ = true;
        ROS_INFO("Calibration results saved to: %s", output_file_.c_str());
    }
    
    void saveCalibrationResults(double rms_error) {
        cv::FileStorage fs(output_file_, cv::FileStorage::WRITE);
        
        if (fs.isOpened()) {
            fs << "camera_matrix" << camera_matrix_;
            fs << "distortion_coefficients" << distortion_coeffs_;
            fs << "image_width" << image_size_.width;
            fs << "image_height" << image_size_.height;
            fs << "rms_reprojection_error" << rms_error;
            fs << "number_of_images" << images_captured_;
            fs << "board_size" << board_size_;
            fs << "square_size" << square_size_;
            
            fs.release();
            
            // 打印标定结果
            ROS_INFO("=== ROS Camera Calibration Results ===");
            ROS_INFO("Camera matrix:");
            for (int i = 0; i < 3; i++) {
                ROS_INFO("  [%.6f, %.6f, %.6f]", 
                        camera_matrix_.at<double>(i, 0),
                        camera_matrix_.at<double>(i, 1),
                        camera_matrix_.at<double>(i, 2));
            }
            ROS_INFO("Distortion coefficients:");
            for (int i = 0; i < distortion_coeffs_.rows; i++) {
                ROS_INFO("  %.6f", distortion_coeffs_.at<double>(i));
            }
            ROS_INFO("RMS error: %.6f pixels", rms_error);
        } else {
            ROS_ERROR("Failed to save calibration results");
        }
    }
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "ros_camera_calibration");
    
    ROSCameraCalibration calibrator;
    
    ros::spin();
    
    return 0;
}
